import jittor as jt
from jittor import nn


class Scale(nn.Module):
    def __init__(self, scale=1.0):
        super().__init__()
        self.scale = jt.Var(scale)

    def execute(self, x):
        return x * self.scale


def NEG_INF_DIAG(n):
    return jt.diag(jt.Var(float('-inf')).repeat(n), 0)


class CrissCrossAttention(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.query_conv = nn.Conv2d(in_channels, in_channels // 8, 1)
        self.key_conv = nn.Conv2d(in_channels, in_channels // 8, 1)
        self.value_conv = nn.Conv2d(in_channels, in_channels, 1)
        self.gamma = Scale(0.)
        self.in_channels = in_channels

    def execute(self, x):
        B, C, H, W = x.size()
        query = self.query_conv(x)
        key = self.key_conv(x)
        value = self.value_conv(x)
        energy_H = jt.linalg.einsum('bchw,bciw->bwhi', query,
                                    key) + NEG_INF_DIAG(H)
        energy_H = energy_H.transpose(1, 2)
        energy_W = jt.linalg.einsum('bchw,bchj->bhwj', query, key)
        attn = nn.softmax(jt.concat([energy_H, energy_W], dim=-1),
                          dim=-1)  # [B,H,W,(H+W)]
        out = jt.linalg.einsum('bciw,bhwi->bchw', value, attn[..., :H])
        out += jt.linalg.einsum('bchj,bhwj->bchw', value, attn[..., H:])

        out = self.gamma(out) + x

        return out

    def __repr__(self):
        s = self.__class__.__name__
        s += f'(in_channels={self.in_channels})'
        return s

class ConvModule(nn.Module):
    _abbr_ = 'conv_block'

    def __init__(self, conv, bn, activate):
        super().__init__()
        self.conv = conv
        self.bn = bn
        self.activate = activate

        # Use msra init by default
        self.init_weights()

    def init_weights(self):
        nn.init.kaiming_normal_(self.conv.weight, mode='fan_out', nonlinearity='relu')
        nn.init.constant_(self.bn.weight, 1)
        nn.init.constant_(self.bn.bias, 0)

    def execute(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.activate(x)
        return x


class CCHead(nn.Module):
    def __init__(self, recurrence=2, **kwargs):
        super(CCHead, self).__init__(**kwargs)
        self.recurrence = recurrence
        self.conv_seg = nn.Conv2d(512, 150, kernel_size=1, stride=1, padding=0)
        self.dropout = nn.Dropout(0.1)
        convs = []
        convs.append(ConvModule(
            conv = nn.Conv2d(2048, 512, kernel_size=3, padding=1, bias=False),
            bn = nn.BatchNorm2d(512),
            activate = nn.ReLU()
        ))
        convs.append(ConvModule(
            conv = nn.Conv2d(512, 512, kernel_size=3, padding=1, bias=False),
            bn = nn.BatchNorm2d(512),
            activate = nn.ReLU()
        ))
        self.convs = nn.Sequential(*convs)
        self.conv_cat = ConvModule(
            conv = nn.Conv2d(2560, 512, kernel_size=3, padding=1, bias=False),
            bn = nn.BatchNorm2d(512),
            activate = nn.ReLU()
        )
        self.cca = CrissCrossAttention(512)

    def execute(self, input, segSize=None):
        """Forward function."""
        x = input
        output = self.convs[0](x)
        for _ in range(self.recurrence):
            output = self.cca(output)
        output = self.convs[1](output)
        output = self.conv_cat(jt.concat([x, output], dim=1))
        output = self.dropout(output)
        output = self.conv_seg(output)

        if segSize is not None:  # inference
            output = nn.interpolate(
                output, size=segSize, mode='bilinear', align_corners=False)
            output = nn.softmax(output, dim=1)
            return output
        
        return output


class FCNHead(nn.Module):
    def __init__(self, **kwargs):
        super(FCNHead, self).__init__(**kwargs)
        self.conv_seg = nn.Conv2d(256, 150, kernel_size=1, stride=1, padding=0)
        self.dropout = nn.Dropout(0.1)
        convs = []
        convs.append(ConvModule(
            conv = nn.Conv2d(1024, 256, kernel_size=3, padding=1, bias=False),
            bn = nn.BatchNorm2d(256),
            activate = nn.ReLU()
        ))
        self.convs = nn.Sequential(*convs)

    def execute(self, input):
        """Forward function."""
        x = input
        x = self.convs[0](x)
        x = self.dropout(x)
        x = self.conv_seg(x)
        
        return x